import torch
from torch.linalg import eig


class PCA():
    def __init__(self, n_component: int = 16, device='cpu') -> None:
        """主成分分析

        Args:
            n_component (int): 保留的主成分数
        """
        super().__init__()
        self.n_component = n_component
        self.device = device

    def CHECK_SHAPE(self, shape: torch.Size) -> None:
        assert len(shape) >= 2, 'Shape of input is expected bigger than 2!!!'
        limit = 1
        for i in range(1, len(shape)):
            limit *= shape[i]
        assert limit >= self.n_component, f'n_component = {self.n_component}, expected <= {limit}'

    @torch.no_grad()
    def fit(self, X: torch.Tensor) -> None:
        """提取主成分

        Args:
            X (torch.Tensor): 待进行主成分分析的输入张量，形状应当为 (batch_size, ...)
        """
        self.CHECK_SHAPE(X.shape)
        Y = X.reshape(X.shape[0], -1).to(self.device)
        self.mean = Y.mean(0)
        Z = Y - self.mean
        
        covariance = Z.T @ Z
        _, eig_vec = eig(covariance)

        self.components = eig_vec[:, :self.n_component]

    @torch.no_grad()
    def transform(self, X: torch.Tensor) -> torch.Tensor:
        """数据降维

        Args:
            X (torch.Tensor): 待降维数据，形状应当为 (batch_size, ...)

        Returns:
            torch.Tensor: 降维后数据
        """
        self.CHECK_SHAPE(X.shape)
        Z = X.reshape(X.shape[0], -1).to(self.device)

        return (Z - self.mean) @ self.components.real
    
    @torch.no_grad()
    def reconstruct(self, X: torch.Tensor) -> torch.Tensor:
        """高维数据重建

        Args:
            X (torch.Tensor): 待重建数据，形状应当为 (batch_size, ...)

        Returns:
            torch.Tensor: 重建后数据
        """
        assert len(X.shape) == 2, 'Shape of input is expected to equal to 2!!!'

        return (X @ self.components.real.T) + self.mean
